package com.atwy.graph.mst;

import java.util.Comparator;
import java.util.List;
import java.util.PriorityQueue;

/**
 * Prim最小生成树:
 * 首先，Prim 算法也使用贪心思想来让生成树的权重尽可能小，也就是「切分定理」。
 * 其次，Prim 算法使用 BFS 算法思想 和 visited 布尔数组避免成环，来保证选出来的边最终形成的一定是一棵树。
 *
 */
public class Prim {
    // 使用邻接表表示图，graph[s] 表示所有与s节点相连的节点及权重，int[]{from,to,weight}，这里的int[] 也可以用Edge类来代替，表示一条边
    private List<int[]>[] graph;
    // 核心数据结构，存储「横切边」的优先级队列
    private PriorityQueue<int[]> pq;
    // 类似 visited 数组的作用，记录哪些节点已经成为最小生成树的一部分
    private boolean[] inMST;
    // 记录最小生成树的权重和
    private int weightSum = 0;

    public Prim(List<int[]>[] graph) {
        this.graph = graph;
        // 比较权重的优先队列
        pq = new PriorityQueue<>(Comparator.comparingInt(a -> a[2]));
        int n = graph.length;
        this.inMST = new boolean[n];

        // 随便选取一个节点，从这个点开始进行切分,比如选取节点0
        int s = 0;
        inMST[s] = true;
        cut(s);
        // 不断切分，向最小生成树中添加边
        while (!pq.isEmpty()){
            int[] edge = pq.poll();
            int to = edge[1];
            int weight = edge[2];
            if(inMST[to]){
                continue;
            }
            this.weightSum += weight;
            inMST[to] = true;
            cut(to);
        }

    }

    private void cut(int s){
        // 遍历s所有相邻的边
        for (int[] edge : this.graph[s]){
            int to = edge[1];
            if(inMST[to]){
                // 说明该边在最小生成树中
                continue;
            }
            pq.offer(edge);
        }
    }

    public int weightSum(){
        return this.weightSum;
    }

    public boolean allConnected(){
        for (int i = 0; i < this.inMST.length; i++) {
            if(!inMST[i]){
                return false;
            }
        }
        return true;
    }




    class Edge{
        private int from;
        private int to;
        private int weight;

        public Edge(int from, int to, int weight) {
            this.from = from;
            this.to = to;
            this.weight = weight;
        }

        public int getFrom() {
            return from;
        }

        public void setFrom(int from) {
            this.from = from;
        }

        public int getTo() {
            return to;
        }

        public void setTo(int to) {
            this.to = to;
        }

        public int getWeight() {
            return weight;
        }

        public void setWeight(int weight) {
            this.weight = weight;
        }
    }
}
